from settings import *
import geopandas as geopd
from utils.postgis import geom_to_ewkt
from shapely.geometry.base import BaseGeometry
from shapely.geometry import LineString


class Table:
    pk = PK_COL
    geom = GEOM_COL
    srid = SRID
    name = ""

    def __init__(self, name):
        self.name = name
        self.dtype_dict = {self.pk: "serial"}

    def drop_table(self):
        with conn:
            with conn.cursor() as cursor:
                cursor.execute(f"drop table {self.name}")

    def create_table(self, drop=False):
        with conn:
            with conn.cursor() as cursor:
                if drop:
                    sql = f"drop table if exists {self.name}"
                    cursor.execute(sql)

                cols = []
                for col, dtype in self.dtype_dict.items():
                    cols.append(f"{col} {dtype}")
                sql = f"create table if not exists {self.name}({','.join(cols)})"
                cursor.execute(sql)

    def update_attr(self, id_col, id_val, update_dict, auto_commit=True):
        with conn.cursor() as cursor:
            sql_set = "set "
            for col, val in update_dict.items():
                if col == self.geom:
                    sql_set += f"{self.geom}={geom_to_ewkt(val, self.srid)},"
                    if isinstance(val, LineString):
                        sql_set += f"{self.length}={val.length},"
                else:
                    if val is None:
                        val = "null"
                    elif isinstance(val, str):
                        val = f"'{val}'"
                    sql_set += f"{col}={val},"
            sql_set = sql_set[:-1]
            if isinstance(id_val, str):
                id_val = f"'{id_val}'"
            sql_str = f"""
                update {self.name}
                {sql_set}
                where {id_col} = {id_val}
            """
            cursor.execute(sql_str)
        if auto_commit:
            conn.commit()

    def select_by_attr(self, col, val, select_cols=None):
        if val is None:
            val = "null"
        elif isinstance(val, str):
            val = f"'{val}'"
        if select_cols is None:  # 返回所有字段
            select_cols_str = "*"
            use_geom = True
        else:
            select_cols_str = ",".join(select_cols)
            if self.geom in select_cols:
                use_geom = True
            else:
                use_geom = False
        sql_str = f"select {select_cols_str} from {self.name} where {col}={val}"
        if use_geom:
            return geopd.read_postgis(sql_str, conn, geom_col=self.geom)
        else:
            return pd.read_sql(sql_str, conn)

    def add_column(self, col_name, set_null=False):
        with conn:
            with conn.cursor() as cursor:
                sql_str = f"alter table {self.name} add column if not exists {col_name} {self.dtype_dict[col_name]}"
                cursor.execute(sql_str)
                if set_null:
                    sql_str = f"update {self.name} set {col_name}=null"
                    cursor.execute(sql_str)

    def select_max_pk(self):
        sql_str = f"select max({self.pk}) from {self.name}"
        pk = pd.read_sql(sql_str, conn).iloc[0, 0]
        if pk is None:
            return 0
        else:
            return pk

    def insert(self, attr_dict):
        with conn:
            with conn.cursor() as cursor:
                sql_cols = []
                sql_vals = []
                if attr_dict:
                    for col, val in attr_dict.items():
                        if val is None:
                            val = "null"
                        elif isinstance(val, str):
                            val = f"'{val}'"
                        elif isinstance(val, BaseGeometry):
                            if isinstance(val, LineString):
                                sql_cols.append(self.length)
                                sql_vals.append(val.length)
                            val = f"st_geomFromEWKT('SRID={self.srid};{val.wkt}')"
                        sql_cols.append(col)
                        sql_vals.append(val)
                sql_str = f"""
                insert into {self.name} ({",".join(sql_cols)})
                values ({",".join([str(x) for x in sql_vals])}) returning {self.pk}
                """
                cursor.execute(sql_str)
                pk_val = cursor.fetchone()[0]
        return pk_val

    def buffer_intersect_by_pt(self, pt_coord, buff_dist, and_where=""):
        sql_str = f"""
        select * from {self.name} where
        st_dwithin(
            st_SetSrid(
                st_Point({pt_coord[0]},{pt_coord[1]}),
                {self.srid}
            ),
            {self.geom},
            {buff_dist}
        )
        {and_where}
        """
        gdf = geopd.read_postgis(sql_str, conn, geom_col=self.geom)
        return gdf

    def buffer_intersect_by_polygon(self, polygon):
        sql_str = f"""
        select * from {self.name} where
        st_intersects(
            {self.geom},
            {geom_to_ewkt(polygon, self.srid)}
        )
        """
        return geopd.read_postgis(sql_str, conn, geom_col=self.geom)

    def buffer_nodes(self, pt, buff_dist, and_where=""):
        sql_str = f"""
        select * from {self.name}
        where
        st_intersects(
            st_buffer(
                {geom_to_ewkt(pt, self.srid)},
                {buff_dist}
            ),
            st_boundary({self.geom})
        )
        {and_where}
    """
        gdf = geopd.read_postgis(sql_str, conn, geom_col=self.geom)
        return gdf

    def delete_by_attr(self, col, val, auto_commit=False):
        with conn.cursor() as cursor:
            if isinstance(val, str):
                val = f"'{val}'"
            sql_str = f"delete from {self.name} where {col}={val}"
            cursor.execute(sql_str)
            if auto_commit:
                conn.commit()

    def add_primary_key(self):
        with conn:
            with conn.cursor() as cursor:
                sql_str = f"alter table {self.name} add primary key ({self.pk})"
                cursor.execute(sql_str)

    def move_rows(self, new_table, where=""):
        with conn:
            with conn.cursor() as cursor:
                sql_str = f"""
                        begin;
                        drop table if exists {new_table.name};
                        select * into {new_table.name} from {self.name} {where};
                        delete from {self.name} {where};
                        commit;
                    """
                cursor.execute(sql_str)

    def vacuum_analyze(self):
        """
        不知道有没有用
        https://blog.csdn.net/qq_35732147/article/details/86212840
        VACUUM: 清理，为了保证数据库的高效运行。每当创建新索引或对表大量更新、插入或删除后，都必须执行清理（VACUUMing）。
        VACUUM命令要求PostgreSQL回收表页面中因记录的更新或删除而留下的任何未使用的空间。
        ANALYZE: 表中加载和删除大容量数据后运行ANALYZE命令。这将强制统计系统收集所有索引列的统计信息。
        """
        with conn:
            with conn.cursor() as cursor:
                old_level = conn.isolation_level
                conn.set_isolation_level(0)
                sql_str = f"VACUUM ANALYZE {self.name}"
                cursor.execute(sql_str)
                conn.set_isolation_level(old_level)

    def create_backup(self, suffix):
        new_table_name = f"{self.name}_{suffix}"
        with conn:
            with conn.cursor() as cursor:
                sql_str = f"""
                                drop table if exists {new_table_name};
                                select * into {new_table_name} from {self.name};
                                alter table {new_table_name} add primary key ({self.pk})"""
                cursor.execute(sql_str)

    def restore_table_from(self, from_table_name):
        max_pk = Table(from_table_name).select_max_pk()
        seq_name = f"{self.name}_{self.pk}_seq"
        with conn:
            with conn.cursor() as cursor:
                sql = f"""
                    drop table if exists {self.name};
                    create table {self.name} as (select * from {from_table_name});
                    create sequence if not exists {seq_name};
                    alter sequence {seq_name} owned by {self.name}.{self.pk};
                    alter sequence {seq_name} restart with {max_pk+1};
                    alter table {self.name} alter column {self.pk} set default nextval('{seq_name}');
                    alter table {self.name} add primary key ({self.pk})"""
                cursor.execute(sql)
                self.create_index(self.geom, True)

    def buffer_intersect_by_layer(self, target_layer):
        sql_str = f"""
            select * from {self.name}
            where 
            st_intersects({self.geom}, 
                (select st_setsrid(
                    st_extent({target_layer.geom}),
                    {self.srid}) 
                    from {target_layer.name})
                )
        """
        gdf = geopd.read_postgis(sql_str, conn, geom_col=self.geom)
        return gdf

    def create_index(self, col, is_geom):
        if is_geom:
            using = "using GIST"
        else:
            using = ""
        with conn:
            with conn.cursor() as cursor:
                sql_str = f"create index on {self.name} {using} ({col})"
                cursor.execute(sql_str)

    def get_pk_list(self):
        df = pd.read_sql(f"select {self.pk} from {self.name} order by {self.pk} asc", conn)
        return df[self.pk].values.tolist()


class NodeTable(Table):
    superior_class = "superior_class"  # 依赖路段的等级，若无依赖道路则为null

    def __init__(self, name):
        Table.__init__(self, name)
        self.dtype_dict = Table("").dtype_dict.copy()
        self.dtype_dict[self.geom] = f"geometry(POINT,{SRID})"
        self.dtype_dict[self.superior_class] = "char"


class RSNodeTable(NodeTable):
    pass


class RoadTable(Table):
    length = LENGTH_COL

    def __init__(self, name):
        Table.__init__(self, name)
        self.dtype_dict = Table("").dtype_dict.copy()
        self.dtype_dict[self.length] = "float"
        self.dtype_dict[self.geom] = f"geometry(linestring, {self.srid})"

    def update_length(self):
        with conn:
            with conn.cursor() as cursor:
                sql_str = f"""
                    alter table {self.name} add column if not exists {self.length} {self.dtype_dict[self.length]};
                    update {self.name} set {self.length}=st_length({self.geom})"""
                cursor.execute(sql_str)

    def simplify(self, tolerance):
        with conn:
            with conn.cursor() as cursor:
                sql = f"""update {self.name} set {self.geom} = st_simplify({self.geom}, {tolerance});
                            update {self.name} set {self.length}=st_length({self.geom});"""
                cursor.execute(sql)


class ReportRoadTable(RoadTable):
    crowid = "crowid"
    segment_code = "ldbm"  # 路段编码
    road_code = "lxbm"  # 路线编码
    road_name = "lxmc"
    ldxlh = "ldxlh"  # 路段序列号
    xzqh = "xzqh"
    start_name = "qdmc"  # 起点名称
    start_stake = "qdzh"  # 起点桩号
    end_name = "zdmc"  # 终点名称
    end_stake = "zdzh"  # 终点桩号
    ldjsdj = "ldjsdj"  # 路段建设等级？
    ldlmlx = "ldlmlx"
    ldlc = "ldlc"
    sfwcfld = 'sfwcfld'
    sfdtl = "sfdtl"
    tbdwdm = "tbdwdm"
    lxbh = "lxbh"  # 路线编号
    lxdjhz = "lxdjhz"

    def __init__(self, name):
        RoadTable.__init__(self, name)
        self.dtype_dict = RoadTable("").dtype_dict.copy()
        self.dtype_dict[self.crowid] = "varchar(36)"
        self.dtype_dict[self.segment_code] = "varchar(20)"
        self.dtype_dict[self.road_code] = "varchar(20)"
        self.dtype_dict[self.road_name] = "varchar(100)"
        self.dtype_dict[self.ldxlh] = "varchar(3)"
        self.dtype_dict[self.xzqh] = "varchar(12)"
        self.dtype_dict[self.start_name] = "varchar(50)"
        self.dtype_dict[self.start_stake] = "float8"
        self.dtype_dict[self.end_name] = "varchar(50)"
        self.dtype_dict[self.end_stake] = "float8"
        self.dtype_dict[self.ldjsdj] = "varchar(1)"
        self.dtype_dict[self.ldlmlx] = "varchar(3)"
        self.dtype_dict[self.ldlc] = "float8"
        self.dtype_dict[self.sfwcfld] = "varchar(1)"
        self.dtype_dict[self.sfdtl] = "varchar(1)"
        self.dtype_dict[self.tbdwdm] = "varchar(20)"
        self.dtype_dict[self.lxbh] = "varchar(6)"
        self.dtype_dict[self.lxdjhz] = "varchar(3)"
        self.dtype_dict[self.geom] = f"geometry(LineString, {self.srid})"


class RemoteSensingRoadTable(RoadTable):
    length = LENGTH_COL
    match_road = MATCH_COL
    buffer_similarity = "buffer_similarity"
    dir_similarity = "dir_similarity"
    hausdorff_dist = "hausdorff_dist"
    cross_num = "cross_num"

    def __init__(self, name):
        RoadTable.__init__(self, name)
        self.dtype_dict = RoadTable("").dtype_dict.copy()
        self.dtype_dict[self.length] = "float"
        self.dtype_dict[self.match_road] = "integer"
        self.dtype_dict[self.buffer_similarity] = "float"
        self.dtype_dict[self.dir_similarity] = "float"
        self.dtype_dict[self.hausdorff_dist] = "float"
        self.dtype_dict[self.cross_num] = "smallint"


class FusedRoadTable(ReportRoadTable, RemoteSensingRoadTable):
    snodeid = SNODEID_COL
    enodeid = ENODEID_COL
    stake_rule = "stake_rule"

    def __init__(self, name):
        super(FusedRoadTable, self).__init__(name)
        self.dtype_dict = ReportRoadTable("").dtype_dict.copy()
        self.dtype_dict.update(RemoteSensingRoadTable("").dtype_dict)
        self.dtype_dict[self.snodeid] = "integer"
        self.dtype_dict[self.enodeid] = "integer"
        self.dtype_dict[self.stake_rule] = "varchar(20)"


class Similarity:
    def __init__(self, pk, rs_geom, buffer_similarity, dir_similarty, h_dist, cross_num):
        self.pk = pk
        self.geom = rs_geom
        self.buffer_similarity = buffer_similarity
        self.dir_similarity = dir_similarty
        self.hausdorff_dist = h_dist
        self.cross_num = cross_num

    @staticmethod
    def similarity_attrs():
        return RemoteSensingRoadTable.buffer_similarity, RemoteSensingRoadTable.dir_similarity, \
               RemoteSensingRoadTable.hausdorff_dist, RemoteSensingRoadTable.cross_num

    def similarity_to_dict(self):
        return {RemoteSensingRoadTable.match_road: self.pk,
                RemoteSensingRoadTable.buffer_similarity: self.buffer_similarity,
                RemoteSensingRoadTable.dir_similarity: self.dir_similarity,
                RemoteSensingRoadTable.hausdorff_dist: self.hausdorff_dist,
                RemoteSensingRoadTable.cross_num: self.cross_num}
